package Server.server.ratelimit.impl;

import Server.server.ratelimit.RateLimit;

public class TokenBucketRateLimitImpl implements RateLimit {

    // 令牌产生速率（ms）
    private static int RATE;
    // 桶容量
    private static int CAPACITY;
    // 当前桶容量
    private volatile int CURRENT_CAPACITY;
    // 上次请求的时间戳
    private volatile long LAST_TIME = System.currentTimeMillis();

    public TokenBucketRateLimitImpl(int rate, int capacity) {
        RATE = rate;
        CAPACITY = capacity;
        CURRENT_CAPACITY = capacity;
    }

    @Override
    public synchronized boolean getToken() {
        // 如果桶内有令牌，直接消费1个令牌并返回True
        if (CURRENT_CAPACITY > 0) {
            CURRENT_CAPACITY--;
            return true;
        }
        // 如果桶内没有令牌，开始计算生成令牌情况
        long currentTime = System.currentTimeMillis();

        // 判断自上次获取令牌是否过去足够时间
        if (currentTime - LAST_TIME >= RATE) {
            // 计算经过时间需要生成读诵阿红令牌
            if ((currentTime - LAST_TIME) / RATE >= 2) {
                CURRENT_CAPACITY += (int) ((currentTime - LAST_TIME) / RATE) - 1;
            }
            // 保证令牌桶容量不会超过最大值 CAPACITY
            if (CURRENT_CAPACITY > CAPACITY) {
                CURRENT_CAPACITY = CAPACITY;
            }
            // 更新时间戳
            LAST_TIME = currentTime;
            return true;
        }
        return false;
    }
}
